package com.dfs.parent.gateway.loadbalancer;

/**
 * @Description TODO
 * @Author dahong
 * @Date 2025/8/7 03:25
 */

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.util.StrUtil;
import com.dfs.parent.common.core.constant.GrayConstant;
import com.dfs.parent.gateway.props.GrayProperties;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.loadbalancer.DefaultRequest;
import org.springframework.cloud.client.loadbalancer.Request;
import org.springframework.cloud.client.loadbalancer.RequestData;
import org.springframework.cloud.client.loadbalancer.RequestDataContext;
import org.springframework.cloud.loadbalancer.core.ServiceInstanceListSupplier; // 正确导入
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.stream.Collectors;
@Slf4j
public class GrayInstanceSupplier implements ServiceInstanceListSupplier {
    private final GrayProperties grayProperties;
    private final ServiceInstanceListSupplier delegate;

    public GrayInstanceSupplier(ServiceInstanceListSupplier delegate,GrayProperties grayProperties) {
        this.delegate = delegate;
        this.grayProperties = grayProperties;
    }

    @Override
    public String getServiceId() {
        return delegate.getServiceId();
    }
    @Override
    public Flux<List<ServiceInstance>> get() {
        return get(new DefaultRequest<>());
    }

    @Override
    public Flux<List<ServiceInstance>> get(Request request) {
        // 1. 检查灰度配置是否启用
        if (!grayProperties.getEnabled()) {
            return delegate.get(request);
        }
        // 如果是灰度请求，过滤实例
        boolean isGrayHeaderTag = hasGrayHeaderTag(request);
        return delegate.get(request)
                    .map(instances -> getInstances(instances,isGrayHeaderTag));
    }
    private boolean hasGrayHeaderTag(Request request) {
        if (!(request.getContext() instanceof RequestDataContext)) {
            return false;
        }
        RequestDataContext context = (RequestDataContext) request.getContext();
        RequestData clientRequest = context.getClientRequest();
        if (clientRequest == null || clientRequest.getHeaders() == null) {
            return false;
        }
        RequestData clientReq = ((RequestDataContext) request.getContext()).getClientRequest();
        return GrayConstant.GRAY_HEADER_TAG_VALUE_TRUE.equals(clientReq.getHeaders().getFirst(GrayConstant.GRAY_HEADER_TAG));

    }

    private List<ServiceInstance> getInstances(List<ServiceInstance> instances,boolean isGrayHeaderTag) {
        if (isGrayHeaderTag) {
            List<ServiceInstance> grayInstances = instances.stream()
                    .filter(instance -> StrUtil.isNotEmpty(instance.getMetadata().get(GrayConstant.GRAY_TAG))
                            && GrayConstant.GRAY_TAG_VALUE_TRUE.equals(instance.getMetadata().get(GrayConstant.GRAY_TAG)))
                    .collect(Collectors.toList());
            // 如果没有灰度实例，返回所有实例
            return CollUtil.isEmpty(grayInstances) ? instances : grayInstances;
        } else {
            List<ServiceInstance> generalInstances = instances.stream()
                    .filter(instance -> StrUtil.isEmpty(instance.getMetadata().get(GrayConstant.GRAY_TAG))
                            || GrayConstant.GRAY_TAG_VALUE_FALSE.equals(instance.getMetadata().get(GrayConstant.GRAY_TAG)))
                    .collect(Collectors.toList());
            // 如果没有非灰度实例，返回所有实例
            return CollUtil.isEmpty(generalInstances) ? instances : generalInstances;
        }
    }
}